#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""YOLO目标检测引擎实现

基于Ultralytics YOLO的高性能目标检测引擎
支持多场景检测：交通、头盔、人脸等
实现场景自适应检测和优化的模型管理
"""

import os
import cv2
import numpy as np
import logging
import torch
import time
from typing import List, Dict, Optional, Tuple, Any
from threading import Lock
from datetime import datetime

from core.detection.base_engine import BaseDetectionEngine
from config import get_config

# 导入头盔检测器
from core.detection.helmet_detector import HelmetDetector

# 设置日志记录器
logger = logging.getLogger(__name__)

class YOLODetectionEngine(BaseDetectionEngine):
    """基于YOLO的目标检测引擎"""
    
    def __init__(self, model_path: Optional[str] = None, device: str = 'auto', **kwargs):
        """初始化YOLO检测引擎
        
        Args:
            model_path: 模型文件路径
            device: 运行设备，'auto'、'cuda'、'cpu'
            **kwargs: 其他配置参数
        """
        super().__init__(device=device, **kwargs)
        self._initialized = False
        
        # 模型相关属性
        self.model = None
        self.model_path = model_path or get_config('detection.model_path', 'models/yolov8n.pt')
        
        # 检测参数
        self.confidence_threshold = get_config('detection.confidence_threshold', 0.5)
        self.iou_threshold = get_config('detection.iou_threshold', 0.45)
        self.max_det = get_config('detection.max_det', 1000)
        
        # 类别信息
        self.classes = []  # 将在模型加载后从模型获取
        self.class_names = {}  # 将在模型加载后从模型获取
        
        # 性能优化设置
        self.fp16 = get_config('performance.fp16', True)
        self.agnostic_nms = get_config('detection.agnostic_nms', False)
        self.fast_inference = get_config('performance.fast_inference', False)
        
        # 线程安全锁
        self.lock = Lock()
        
        # 检测模式 (traffic, helmet, face, all)
        self.detection_mode = kwargs.get('detection_mode', 'traffic')
        
        # 场景特定参数配置
        self.scene_params = {
            'traffic': {
                'confidence_threshold': 0.3,
                'target_classes': [0, 1, 2, 3, 5, 7]  # person, bicycle, car, motorbike, bus, truck
            },
            'helmet': {
                'confidence_threshold': 0.3,
                'target_classes': [0, 3]  # person, motorbike
            },
            'face': {
                'confidence_threshold': 0.4,
                'target_classes': None  # 使用默认类别或专门的人脸检测模型
            },
            'all': {
                'confidence_threshold': 0.3,
                'target_classes': None  # 检测所有类别
            }
        }
        
        # 头盔检测器
        self.helmet_detector = None
        
        # 场景模型路径映射
        self.scene_model_paths = {
            'traffic': self.model_path,
            'helmet': self.model_path,
            'face': self.model_path,
            'all': self.model_path
        }
        
        # 上次检测结果缓存
        self.last_detections = None
        
        # 尝试自动初始化
        self.initialize()
    
    def initialize(self) -> bool:
        """初始化检测引擎
        
        Returns:
            bool: 是否初始化成功
        """
        try:
            logger.info(f"初始化YOLO检测引擎，设备: {self.device}, 模型路径: {self.model_path}")
            
            # 确保模型路径存在
            if not os.path.isabs(self.model_path):
                base_dir = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
                self.model_path = os.path.join(base_dir, self.model_path)
            
            # 检查模型文件是否存在，如果不存在则使用默认的YOLOv8n模型
            if not os.path.exists(self.model_path):
                logger.warning(f"模型文件不存在: {self.model_path}, 将使用默认的YOLOv8n模型")
                # 尝试从ultralytics下载默认模型
                try:
                    from ultralytics import YOLO
                    self.model = YOLO('yolov8n.pt')
                    logger.info("已成功加载默认YOLOv8n模型")
                    
                    # 从模型获取类别信息
                    self.classes = list(range(len(self.model.names)))
                    self.class_names = self.model.names
                    logger.info(f"从模型获取到 {len(self.class_names)} 个类别")
                except Exception as e:
                    logger.error(f"无法加载默认模型: {str(e)}")
                    self._initialized = False
                    return False
            else:
                # 加载指定的模型
                from ultralytics import YOLO
                self.model = YOLO(self.model_path)
                logger.info(f"成功加载模型: {self.model_path}")
                
                # 从模型获取类别信息
                self.classes = list(range(len(self.model.names)))
                self.class_names = self.model.names
                logger.info(f"从模型获取到 {len(self.class_names)} 个类别")
            
            # 设置设备
            if self.device == 'auto':
                # 自动选择设备
                self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
                logger.info(f"自动选择设备: {self.device}")
            
            # 验证设备是否可用
            if self.device == 'cuda' and not torch.cuda.is_available():
                logger.warning("CUDA不可用，切换到CPU")
                self.device = 'cpu'
                self.fp16 = False  # CPU不支持FP16
            
            # 获取模型内置的类别名称
            if hasattr(self.model, 'names') and isinstance(self.model.names, dict):
                self.class_names.update(self.model.names)
            
            self._initialized = True
            logger.info("YOLO检测引擎初始化成功")
            
            # 执行优化
            self.optimize()
            
            # 初始化头盔检测器
            try:
                self.helmet_detector = HelmetDetector()
                logger.info("头盔检测器初始化成功")
            except Exception as e:
                logger.warning(f"头盔检测器初始化失败: {str(e)}")
                logger.warning("建议为头盔检测训练专用模型以获得最佳效果")
                self.helmet_detector = None
            
            return True
            
        except ImportError as e:
            logger.error(f"缺少必要的库: {str(e)}")
            return False
        except Exception as e:
            logger.error(f"YOLO检测引擎初始化失败: {str(e)}")
            return False
    
    def is_healthy(self) -> bool:
        """
        检查引擎是否健康运行
        
        Returns:
            bool: 引擎是否健康
        """
        return self._initialized and self.model is not None
    
    def detect(self, frame: np.ndarray, **kwargs) -> List[Dict[str, Any]]:
        """执行目标检测
        
        Args:
            frame: 输入图像帧
            **kwargs: 额外检测参数
                - confidence_threshold: 置信度阈值
                - classes: 要检测的类别列表
                - agnostic_nms: 是否使用类别无关的NMS
                - scene_type: 指定检测场景类型（可选）
                - enable_tracking: 是否启用目标跟踪
                
        Returns:
            List[Dict]: 检测结果列表，包含场景特定标记和分析信息
        """
        if not self._initialized or self.model is None:
            logger.error("检测引擎未初始化，无法执行检测")
            return []
        
        start_time = time.time()
        
        try:
            # 验证输入帧
            if frame is None or not isinstance(frame, np.ndarray) or frame.size == 0:
                logger.error("无效的输入帧")
                return []
            
            # 获取当前场景类型（优先使用传入的参数，否则使用当前模式）
            scene_type = kwargs.get('scene_type', self.detection_mode)
            
            # 场景特定预处理
            processed_frame = self._preprocess_frame(frame.copy(), scene_type)
            
            # 从场景参数配置中获取配置
            scene_config = self.scene_params.get(scene_type, {})
            confidence_threshold = kwargs.get('confidence_threshold', 
                                            scene_config.get('confidence_threshold', self.confidence_threshold))
            target_classes = scene_config.get('target_classes', None)
            classes = kwargs.get('classes', target_classes)
            agnostic_nms = kwargs.get('agnostic_nms', self.agnostic_nms)
            
            # 场景特定参数调整
            if scene_type == 'traffic':
                # 交通场景检测参数优化
                if classes is None:
                    classes = [0, 1, 2, 3, 5, 7]  # person, bicycle, car, motorbike, bus, truck
                # 交通场景IOU阈值可以适当降低以分离密集车辆
                current_iou = self.iou_threshold
                if current_iou > 0.4:
                    current_iou = 0.4
            elif scene_type == 'helmet':
                # 头盔检测参数优化
                if confidence_threshold > 0.3:
                    confidence_threshold = 0.3  # 降低阈值以提高召回率
                # 针对头盔场景，确保检测人或摩托车
                if classes is None:
                    classes = [0, 3]  # person, motorbike
            elif scene_type == 'face':
                # 人脸检测参数优化
                if confidence_threshold > 0.4:
                    confidence_threshold = 0.4  # 降低阈值以提高召回率
            elif scene_type == 'all':
                # 全场景模式，检测所有类别
                classes = None
                # 全场景模式需要更精确的检测
                if confidence_threshold < 0.4:
                    confidence_threshold = 0.4
            
            with self.lock:
                # 执行检测
                results = self.model(
                    processed_frame,
                    conf=confidence_threshold,
                    iou=self.iou_threshold if scene_type != 'traffic' else current_iou,
                    max_det=self.max_det,
                    classes=classes,
                    device=self.device,
                    half=self.fp16,
                    agnostic_nms=agnostic_nms,
                    verbose=False
                )
            
            # 处理检测结果
            detections = []
            for result in results:
                boxes = result.boxes
                for box in boxes:
                    # 获取边界框坐标
                    x1, y1, x2, y2 = map(int, box.xyxy[0].tolist())
                    
                    # 获取置信度和类别ID
                    confidence = float(box.conf[0])
                    class_id = int(box.cls[0])
                    
                    # 获取类别名称
                    class_name = self.class_names.get(class_id, f'class_{class_id}')
                    
                    # 计算中心点
                    center_x = (x1 + x2) // 2
                    center_y = (y1 + y2) // 2
                    
                    # 创建基本检测结果，增加场景特定标记
                    detection = {
                        'bbox': [x1, y1, x2, y2],
                        'x1': x1,
                        'y1': y1,
                        'x2': x2,
                        'y2': y2,
                        'center': (center_x, center_y),
                        'confidence': confidence,
                        'class_id': class_id,
                        'class_name': class_name,
                        'area': (x2 - x1) * (y2 - y1),
                        'has_helmet': None,  # 头盔检测标记
                        'helmet_confidence': None,  # 头盔检测置信度
                        'scene_type': scene_type,  # 当前检测场景
                        'tracking_id': None,  # 用于目标跟踪
                        'scene_specific_info': {}  # 场景特定信息
                    }
                    
                    # 场景特定标记
                    if scene_type == 'traffic':
                        # 交通场景标记
                        if class_id in [0, 1, 2, 3, 5, 7]:
                            detection['is_traffic_target'] = True
                            # 标记为车辆或行人
                            if class_id in [1, 2, 3, 5, 7]:  # 车辆类别
                                detection['scene_specific_info']['is_vehicle'] = True
                                detection['scene_specific_info']['vehicle_type'] = class_name
                            elif class_id == 0:  # 行人
                                detection['scene_specific_info']['is_pedestrian'] = True
                        else:
                            detection['is_traffic_target'] = False
                    elif scene_type == 'helmet':
                        # 头盔场景标记
                        if class_id in [0, 3]:  # 人或摩托车
                            detection['is_helmet_target'] = True
                            # 针对摩托车骑手的特殊标记
                            if class_id == 0:
                                detection['scene_specific_info']['is_rider_candidate'] = True
                            elif class_id == 3:
                                detection['scene_specific_info']['is_motorbike'] = True
                        else:
                            detection['is_helmet_target'] = False
                    
                    # 添加到检测结果列表
                    detections.append(detection)
            
            # 执行场景特定后处理
            detections = self._postprocess_results(detections, frame.shape[:2], scene_type)
            
            # 头盔检测特殊处理
            if (scene_type == 'helmet' or scene_type == 'all') and self.helmet_detector:
                try:
                    logger.debug("执行头盔检测")
                    detections = self.helmet_detector.process_detections_with_helmets(frame, detections)
                    
                    # 增强头盔检测结果分析
                    for det in detections:
                        if det.get('is_helmet_target', False) and det['class_id'] == 0:
                            # 分析头盔佩戴情况
                            if det.get('has_helmet') is False:
                                logger.warning(f"检测到未佩戴头盔的骑手，置信度: {det.get('helmet_confidence', 'N/A')}")
                                det['scene_specific_info']['violation_type'] = 'no_helmet'
                                det['scene_specific_info']['violation_confidence'] = det.get('helmet_confidence', 1.0)
                except Exception as e:
                    logger.error(f"头盔检测过程发生错误: {str(e)}")
            elif scene_type == 'helmet' and not self.helmet_detector:
                logger.warning("头盔检测器未初始化，建议训练专用模型")
                logger.warning("使用默认YOLO模型进行头盔检测效果可能不理想，推荐以下方案:")
                logger.warning("1. 收集头盔数据集并训练专用头盔检测模型")
                logger.warning("2. 或者使用两阶段检测: 先检测人，再检测头盔")
            
            # 更新性能统计
            self.update_performance_stats(start_time, len(detections))
            
            # 缓存检测结果
            self.last_detections = detections
            
            return detections
            
        except Exception as e:
            logger.error(f"执行检测失败: {str(e)}")
            return []
    
    def _preprocess_frame(self, frame: np.ndarray, scene_type: str) -> np.ndarray:
        """场景特定的图像预处理
        
        Args:
            frame: 输入图像帧
            scene_type: 场景类型
            
        Returns:
            np.ndarray: 预处理后的图像帧
        """
        # 基本预处理
        processed_frame = frame.copy()
        
        # 场景特定预处理
        if scene_type == 'traffic':
            # 交通场景预处理：增强对比度以更好地检测车辆
            alpha = 1.2  # 对比度增益
            beta = 10    # 亮度增益
            processed_frame = cv2.convertScaleAbs(processed_frame, alpha=alpha, beta=beta)
        elif scene_type == 'helmet':
            # 头盔场景预处理：增强顶部区域对比度以更好地检测头盔
            # 创建顶部区域ROI
            height, width = frame.shape[:2]
            top_region = frame[:height//2, :]
            # 增强顶部区域
            top_enhanced = cv2.convertScaleAbs(top_region, alpha=1.3, beta=5)
            # 合并增强后的顶部区域
            processed_frame[:height//2, :] = top_enhanced
        elif scene_type == 'face':
            # 人脸场景预处理：直方图均衡化以改善光线条件
            if len(frame.shape) == 3:
                # 转换为灰度图
                gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
                # 直方图均衡化
                equalized = cv2.equalizeHist(gray)
                # 转回彩色图
                processed_frame = cv2.cvtColor(equalized, cv2.COLOR_GRAY2BGR)
        
        return processed_frame
    
    def _postprocess_results(self, detections: List[Dict[str, Any]], frame_shape: Tuple[int, int], 
                            scene_type: str = None) -> List[Dict[str, Any]]:
        """场景特定的检测结果后处理
        
        Args:
            detections: 原始检测结果
            frame_shape: 图像帧形状 (height, width)
            scene_type: 当前场景类型
            
        Returns:
            List[Dict]: 后处理后的检测结果
        """
        if scene_type is None:
            scene_type = self.detection_mode
            
        height, width = frame_shape
        
        # 过滤无效检测框
        valid_detections = []
        for det in detections:
            # 检查边界框是否有效
            x1, y1, x2, y2 = det['bbox']
            if x1 < 0 or y1 < 0 or x2 > width or y2 > height:
                logger.debug("检测到边界框超出图像范围，已过滤")
                continue
            
            # 检查边界框尺寸
            bbox_width = x2 - x1
            bbox_height = y2 - y1
            
            # 场景特定的边界框尺寸过滤
            if scene_type == 'traffic':
                # 交通场景：允许检测较小的远处车辆
                min_size = 5
            elif scene_type == 'helmet':
                # 头盔场景：需要足够大以看清头部
                min_size = 8
            elif scene_type == 'face':
                # 人脸场景：需要足够大以识别面部
                min_size = 10
            else:
                min_size = 10
            
            if bbox_width < min_size or bbox_height < min_size:
                logger.debug(f"检测到过小的边界框，已过滤（场景: {scene_type}）")
                continue
            
            valid_detections.append(det)
        
        # 场景特定的非极大值抑制（NMS）
        if scene_type == 'traffic' and len(valid_detections) > 0:
            valid_detections = self._apply_nms(valid_detections, iou_threshold=0.35)
        
        # 按置信度排序
        valid_detections.sort(key=lambda x: x['confidence'], reverse=True)
        
        # 场景特定的目标过滤
        if scene_type == 'helmet':
            # 过滤出与头盔检测相关的目标
            helmet_related_detections = []
            for det in valid_detections:
                if det['class_id'] in [0, 3]:  # 人或摩托车
                    helmet_related_detections.append(det)
            
            # 如果有相关目标，使用这些目标；否则保留所有目标
            if helmet_related_detections:
                valid_detections = helmet_related_detections
        elif scene_type == 'all':
            # 全场景模式下，限制返回的检测数量以提高性能
            if len(valid_detections) > 50:
                valid_detections = valid_detections[:50]
        
        return valid_detections
    
    def _apply_nms(self, detections: List[Dict[str, Any]], iou_threshold: float = 0.45) -> List[Dict[str, Any]]:
        """对检测结果应用非极大值抑制
        
        Args:
            detections: 检测结果列表
            iou_threshold: IOU阈值
            
        Returns:
            List[Dict]: NMS处理后的检测结果
        """
        if not detections:
            return []
        
        # 按类别分组
        class_groups = {}
        for det in detections:
            class_id = det['class_id']
            if class_id not in class_groups:
                class_groups[class_id] = []
            class_groups[class_id].append(det)
        
        # 对每个类别应用NMS
        nms_results = []
        for class_id, class_detections in class_groups.items():
            # 提取边界框坐标
            boxes = np.array([det['bbox'] for det in class_detections])
            scores = np.array([det['confidence'] for det in class_detections])
            
            # 应用NMS
            indices = cv2.dnn.NMSBoxes(boxes.tolist(), scores.tolist(), 0.1, iou_threshold)
            
            # 收集NMS后的结果
            if len(indices) > 0:
                for i in indices.flatten():
                    nms_results.append(class_detections[i])
        
        return nms_results
    
    def set_detection_mode(self, mode: str, force_reload: bool = False) -> bool:
        """设置检测模式
        
        Args:
            mode: 检测模式 ('traffic', 'helmet', 'face', 'all')
            force_reload: 是否强制重新加载模型
            
        Returns:
            bool: 是否设置成功
        """
        if mode in ['traffic', 'helmet', 'face', 'all']:
            old_mode = self.detection_mode
            
            # 如果模式相同且不强制重新加载，则不执行切换
            if mode == old_mode and not force_reload:
                logger.debug(f"当前已经是{mode}模式，无需切换")
                return True
                
            self.detection_mode = mode
            logger.info(f"检测模式已切换为: {mode}")
            
            # 调用场景切换回调
            self.on_scene_switch(old_mode, mode)
            
            # 场景特定的模型和参数调整
            if mode == 'helmet':
                # 头盔模式特殊处理
                if not self.helmet_detector:
                    logger.warning("头盔检测模式已启用，但头盔检测器未初始化")
                    logger.warning("使用通用YOLO模型进行头盔检测效果可能不理想")
                    logger.warning("推荐方案:")
                    logger.warning("1. 收集头盔数据集并训练专用头盔检测模型")
                    logger.warning("2. 使用YOLOv8预训练模型进行迁移学习")
                    logger.warning("3. 或者使用两阶段检测方法")
                
                # 尝试动态调整参数以提高头盔检测效果
                logger.info("调整参数以优化头盔检测效果")
                self.confidence_threshold = 0.3
                self.iou_threshold = 0.4
            
            # 尝试加载场景专用模型
            if self.scene_model_paths.get(mode) and self.scene_model_paths[mode] != self.model_path:
                try:
                    self._load_model_for_scene(mode)
                except Exception as e:
                    logger.error(f"加载场景{mode}专用模型失败: {str(e)}")
                    logger.info("继续使用当前模型")
            
            return True
        else:
            logger.error(f"无效的检测模式: {mode}")
            return False
    
    def _load_model_for_scene(self, scene_type: str):
        """为特定场景加载专用模型
        
        Args:
            scene_type: 场景类型
        """
        model_path = self.scene_model_paths.get(scene_type)
        if not model_path or not os.path.exists(model_path):
            logger.warning(f"场景{scene_type}的专用模型不存在: {model_path}")
            return
        
        logger.info(f"加载场景{scene_type}专用模型: {model_path}")
        
        # 保存当前模型路径
        self.model_path = model_path
        
        # 重新初始化引擎
        self.initialize()
        logger.info(f"场景{scene_type}专用模型加载完成")
    
    def on_scene_switch(self, old_scene: str, new_scene: str, camera_id: Optional[str] = None):
        """场景切换时的回调函数
        
        Args:
            old_scene: 切换前的场景
            new_scene: 切换后的场景
            camera_id: 摄像头ID（可选）
        """
        logger.info(f"检测引擎场景切换: {old_scene} -> {new_scene}, 摄像头ID: {camera_id}")
        
        # 清除之前的检测结果缓存
        self.last_detections = None
        
        # 加载场景专用模型
        if new_scene in self.scene_model_paths:
            scene_model_path = self.scene_model_paths[new_scene]
            if scene_model_path and scene_model_path != self.model_path:
                try:
                    logger.info(f"加载{new_scene}场景专用模型: {scene_model_path}")
                    # 这里可以实现动态加载场景专用模型
                    # 为了简单起见，目前使用同一个模型但调整参数
                except Exception as e:
                    logger.error(f"加载场景专用模型失败: {str(e)}")
    
    def load_scene_model(self, scene_type: str, model_path: str) -> bool:
        """为特定场景加载专用模型
        
        Args:
            scene_type: 场景类型
            model_path: 模型路径
            
        Returns:
            bool: 是否成功加载
        """
        if scene_type not in ['traffic', 'helmet', 'face', 'all']:
            logger.error(f"不支持的场景类型: {scene_type}")
            return False
        
        if not os.path.exists(model_path):
            logger.error(f"模型文件不存在: {model_path}")
            return False
        
        try:
            self.scene_model_paths[scene_type] = model_path
            logger.info(f"为{scene_type}场景设置专用模型: {model_path}")
            
            # 如果当前就是该场景，尝试切换模型
            if self.detection_mode == scene_type:
                logger.info(f"当前场景为{scene_type}，建议重启引擎以应用新模型")
            
            return True
        except Exception as e:
            logger.error(f"加载场景模型失败: {str(e)}")
            return False
    
    def set_confidence_threshold(self, threshold: float) -> bool:
        """设置置信度阈值
        
        Args:
            threshold: 置信度阈值 (0.0-1.0)
            
        Returns:
            bool: 是否设置成功
        """
        if 0.0 <= threshold <= 1.0:
            self.confidence_threshold = threshold
            logger.info(f"置信度阈值已设置为: {threshold}")
            return True
        else:
            logger.error("置信度阈值必须在0.0-1.0之间")
            return False
    
    def optimize(self):
        """优化检测引擎性能"""
        try:
            logger.info("优化YOLO检测引擎性能")
            
            # 根据设备类型和配置进行优化
            if self.device == 'cuda':
                # GPU优化
                torch.backends.cudnn.benchmark = True
                logger.info("已启用CuDNN基准测试优化")
                
                # 如果启用快速推理模式
                if self.fast_inference:
                    logger.info("启用快速推理模式")
                    # 可以在这里添加更多GPU优化设置
            
            logger.info("YOLO检测引擎性能优化完成")
            
        except Exception as e:
            logger.error(f"优化失败: {str(e)}")
    
    def get_model_info(self) -> Dict[str, Any]:
        """获取模型信息和场景特定配置
        
        Returns:
            Dict: 包含模型信息、场景配置和优化建议的字典
        """
        if not self._initialized or self.model is None:
            return {}
        
        try:
            model_info = {
                'model_path': self.model_path,
                'model_name': self.model.__class__.__name__ if hasattr(self.model, '__class__') else 'Unknown',
                'device': self.device,
                'detection_mode': self.detection_mode,
                'confidence_threshold': self.confidence_threshold,
                'iou_threshold': self.iou_threshold,
                'max_detections': self.max_det,
                'fp16_enabled': self.fp16,
                'class_count': len(self.class_names),
                'classes': list(self.class_names.items()),
                'helmet_detector_enabled': self.helmet_detector is not None,
                'scene_models': self.scene_model_paths,
                'scene_params': self.scene_params,
                'current_scene_params': self.scene_params.get(self.detection_mode, {}),
                'optimization_suggestions': []
            }
            
            # 添加场景特定的优化建议
            if self.detection_mode == 'helmet':
                model_info['optimization_suggestions'].append("头盔检测场景建议使用专用训练模型")
                model_info['optimization_suggestions'].append("建议收集包含各种光线条件的头盔数据集")
                model_info['optimization_suggestions'].append("考虑使用两阶段检测方法提高准确率")
                
                # 模型训练建议
                model_info['training_recommendations'] = {
                    'dataset': "建议收集至少5000张包含头盔和不带头盔的图像",
                    'augmentation': "使用旋转、缩放、裁剪、亮度变化等数据增强",
                    'transfer_learning': "建议基于YOLOv8预训练模型进行迁移学习",
                    'epochs': "推荐训练100-200个epoch",
                    'batch_size': "根据GPU内存选择合适的批次大小"
                }
            elif self.detection_mode == 'traffic':
                model_info['optimization_suggestions'].append("交通场景建议使用针对车辆训练的模型")
                model_info['optimization_suggestions'].append("考虑添加车辆属性识别（颜色、类型等）")
            
            # 通用优化建议
            model_info['optimization_suggestions'].append("对于实时应用，建议使用YOLOv8n或YOLOv8s模型")
            model_info['optimization_suggestions'].append("定期更新模型以适应新的场景和目标")
            
            return model_info
            
        except Exception as e:
            logger.error(f"获取模型信息失败: {str(e)}")
            return {}
    
    def update_performance_stats(self, start_time: float, detection_count: int):
        """更新性能统计信息
        
        Args:
            start_time: 检测开始时间
            detection_count: 检测到的目标数量
        """
        try:
            elapsed_time = time.time() - start_time
            fps = 1.0 / elapsed_time if elapsed_time > 0 else 0
            
            # 这里可以实现更复杂的性能统计逻辑
            logger.debug(f"检测耗时: {elapsed_time:.3f}秒, FPS: {fps:.1f}, 检测目标数: {detection_count}")
        except Exception as e:
            logger.error(f"更新性能统计失败: {str(e)}")
    
    def release(self):
        """释放检测引擎资源"""
        try:
            with self.lock:
                if self.model is not None:
                    # 释放模型资源
                    if hasattr(self.model, 'to'):
                        # 将模型移至CPU以释放GPU内存
                        self.model.to('cpu')
                    # 清除模型引用
                    self.model = None
            
            # 释放头盔检测器资源
            if self.helmet_detector:
                try:
                    self.helmet_detector.release()
                    self.helmet_detector = None
                    logger.info("头盔检测器资源已释放")
                except Exception as e:
                    logger.error(f"释放头盔检测器资源失败: {str(e)}")
            
            # 清除缓存和临时数据
            self.last_detections = None
            
            # 重置初始化状态
            self._initialized = False
            logger.info("YOLO检测引擎资源已释放")
            
        except Exception as e:
            logger.error(f"释放检测引擎资源失败: {str(e)}")

# 确保time模块已导入（已在文件开头导入）

# 导入必要的模块
import time